import os
from dateutil import parser
import unittest
import numpy as np
import pandas as pd
from investment.lixinger_api import SAVE_INDEX_DIR
from investment.constants import INDUSTRY_CATEGORY, A_INDEX_DICT
from investment.time_series import IndexTimeSeries, StockTimeSeries, IndexBasketTimeSeries, concat_time
from investment.time_series import filter_indexes_connection, filter_indexes_by_spacer


def read_index_csv(csv_file):
    csv = pd.read_csv(csv_file)
    dates = csv['日期']
    prices = csv['收盘点位']
    pe = csv["PE-TTM市值加权"]
    return dates, prices, pe


def csv_date_parser(date):
    return parser.parse(date).strftime("%Y/%#m/%#d")


class testIndexSeries:
    def __init__(self,code, csv_file):
        self.code = code
        self.dates, self.prices, self.pe = read_index_csv(csv_file)

    def check_date(self, time_series):
        assert isinstance(time_series, IndexTimeSeries)
        start_date = time_series.dates_series[0]
        end_date = time_series.dates_series[-1]
        start_date = csv_date_parser(start_date)
        end_date = csv_date_parser(end_date)
        start_idx = np.where(self.dates == start_date)[0][0]
        end_idx = np.where(self.dates == end_date)[0][0]
        date_check = (self.dates[start_idx:end_idx-1:-1] == list(map(csv_date_parser, time_series.dates_series))).all()
        price_check = (self.prices[start_idx:end_idx-1:-1] == time_series.get_fundamental_metrics("CP")).all()
        pe_check = (np.abs(self.pe[start_idx:end_idx-1:-1] - time_series.get_fundamental_metrics("PE")) < 1e-4).all()
        check_all = date_check and price_check and pe_check
        return check_all, self.pe[start_idx:end_idx-1:-1].array


class TestDataClass(unittest.TestCase):
    def test_index(self):
        index = IndexTimeSeries('000300', "2020-03-31", "2021-03-31")
        index.set_json_path(
            os.path.join(SAVE_INDEX_DIR, '000300', "000300_2020-03-31_2021-03-31_index_fundamental.json"),
            os.path.join(SAVE_INDEX_DIR, '000300', "000300_2020-03-31_2021-03-31_index_finance.json"),
            os.path.join(SAVE_INDEX_DIR, '000300', "['000300']_index_samples.json"))
        index.fill_metrics()
        result_1 = index + index
        result_cp = index.deep_copy()
        result_cp += index
        c = 0.5 * index
        test_index = testIndexSeries("000300", r"D:\Work\Projects\Jazz\test\data\沪深300_PE-TTM_市值加权_所有_20221019.csv")
        check, test_pe_ttm = test_index.check_date(index)
        self.assertTrue(check)
        self.assertTrue(id(result_1) != id(index) and id(result_cp) != id(index) and id(c) != id(index))
        self.assertTrue((np.abs(result_cp.get_fundamental_metrics('PE') - test_pe_ttm * 2) < 1e-4).all())
        self.assertTrue((np.abs(result_1.get_fundamental_metrics('PE') - test_pe_ttm * 2) < 1e-4).all())
        self.assertTrue((np.abs(index.get_fundamental_metrics('PE') - test_pe_ttm) < 1e-4).all())
        self.assertTrue((np.abs(c.get_fundamental_metrics('PE') - test_pe_ttm * 0.5 < 1e-4)).all())
        print("pass index test")

    def test_20_years_index(self):
        a_market = IndexTimeSeries(A_INDEX_DICT['沪深300'], "2002-07-01", "2022-07-01", fill_flag=True)
        test_index = testIndexSeries("000300", r"D:\Work\Projects\Jazz\test\data\沪深300_PE-TTM_市值加权_所有_20221019.csv")
        check, test_pe_ttm = test_index.check_date(a_market)
        self.assertTrue(check)
        print("pass 20years hushen300 index pe")

    def test_concat(self):
        test_res = StockTimeSeries('000963', "2011-08-27", "2022-07-21", INDUSTRY_CATEGORY.non_finance, fill_flag=True)
        stock = StockTimeSeries('000963', "2011-08-27", "2021-08-27", INDUSTRY_CATEGORY.non_finance, fill_flag=True)
        stock2 = StockTimeSeries('000963', "2021-08-28", "2022-07-21", INDUSTRY_CATEGORY.non_finance, fill_flag=True)
        stock3 = concat_time(stock, stock2)
        self.assertTrue(test_res == stock3)
        print("pass concat TimeSeries")

    def test_stock(self):
        stock = StockTimeSeries('000963', "2011-08-27", "2021-08-27", INDUSTRY_CATEGORY.non_finance)
        stock.fill_metrics()
        result_1 = stock + stock
        result_cp = stock.deep_copy()
        result_cp += stock
        c = 0.5 * stock
        self.assertTrue(id(result_1) != id(stock) and id(result_cp) != id(stock) and id(c) != id(stock))
        test_b_roe = np.array([0.32063895, 0.30955423, 0.32199326, 0.31829517, 0.30920685,
                               0.3000325, 0.3075846, 0.30271696, 0.2881501, 0.27912304,
                               0.28534021, 0.28521577, 0.32247483, 0.34072717, 0.34394375,
                               0.38463017, 0.4429988, 0.39991925, 0.24963178, 0.27222837,
                               0.27593061, 0.28210998, 0.21500672, 0.23510095, 0.23268377,
                               0.22956047, 0.22785572, 0.24797442, 0.2420313, 0.24905073,
                               0.25423842, 0.25669803, 0.25169984, 0.25119212, 0.24826057,
                               0.23641565, 0.22818985, 0.20749575, 0.16950304, 0.16335265])
        self.assertTrue((np.abs(result_1.get_finance_metrics("ROE") - test_b_roe * 2) < 1e-4).all())
        self.assertTrue((np.abs(result_cp.get_finance_metrics("ROE") - test_b_roe * 2) < 1e-4).all())
        self.assertTrue((np.abs(stock.get_finance_metrics("ROE") - test_b_roe) < 1e-4).all())
        self.assertTrue((np.abs(c.get_finance_metrics("ROE") - test_b_roe * 0.5 < 1e-4)).all())
        print("pass stock test")

    def test_index_basket(self):
        index_300 = IndexTimeSeries('000300', "2011-10-01", "2021-10-01", fill_flag=True)
        index_500 = IndexTimeSeries('000905', "2011-10-01", "2021-10-01", fill_flag=True)
        test_index_boxes = IndexBasketTimeSeries("2011-10-01", "2021-10-01",
                                                 [A_INDEX_DICT['中证500'], A_INDEX_DICT['沪深300']])
        test_pe = (index_300.get_fundamental_metrics("PE") + index_500.get_fundamental_metrics("PE")) / 2
        self.assertTrue((np.all(np.abs(test_index_boxes.basket_time_series.get_fundamental_metrics("PE") - test_pe)
                                < 1e-4)))
        cp_300 = index_300.get_fundamental_metrics("CP")
        cp_500 = index_500.get_fundamental_metrics("CP")
        cp_box = test_index_boxes.get_merged_attribute('get_net_value')
        self.assertTrue((np.abs(cp_300/cp_300[0]*0.5 + 0.5*cp_500/cp_500[0] - cp_box) < 1e-4).all())
        print("pass test index box")

    # test date_utils: search trade day
    def test_search_trade_day(self):
        from date_utilities import search_trade_day, DATE_FORMAT, interp_trading_days
        from dateutil import parser, rrule
        start_day = "2020-03-31"
        end_day = "2021-03-31"
        my_index = IndexTimeSeries('399995', start_day, end_day, fill_flag=True)
        dates = my_index.dates_series
        assert search_trade_day(dates, start_day) == (0, start_day)
        assert search_trade_day(dates, end_day) == (243, end_day)
        result = []
        rule_days = rrule.rrule(rrule.DAILY, dtstart=parser.parse(start_day), until=parser.parse(end_day))
        for day in rule_days:
            result.append(search_trade_day(dates, day.strftime(DATE_FORMAT)))
        cache_index = 0
        # check increasing
        for res in result:
            index = res[0]
            if index < cache_index:
                raise AssertionError("test search trade day failed!")
            cache_index = index

        interp_day = interp_trading_days(dates)
        interp_value = list(interp_day.values())
        for i, res in enumerate(result):
            if res[0] != interp_value[i]:
                print(res, interp_value)
        print("pass date utilities test")

    def test_filter_index_connection(self):
        a = np.array([1, 4, 5, 7])
        b = np.array([1, 4, 5, 7, 10])
        c = np.array([1, 4])
        d = np.array([1, 2])
        e = np.array([1])
        pad = 2
        self.assertTrue((filter_indexes_connection(a, pad) == np.array([1, 5])).all())
        self.assertTrue((filter_indexes_connection(b, pad) == np.array([1, 5, 10])).all())
        self.assertTrue((filter_indexes_connection(c, pad) == np.array([1, 4])).all())
        self.assertTrue((filter_indexes_connection(d, pad) == np.array([1])).all())
        self.assertTrue((filter_indexes_connection(e, pad) == np.array([1])).all())

        self.assertTrue((filter_indexes_by_spacer(a, pad) == np.array([1, 4, 7])).all())
        self.assertTrue((filter_indexes_by_spacer(b, 3) == np.array([1, 5, 10])).all())
        self.assertTrue((filter_indexes_by_spacer(c, pad) == np.array([1, 4])).all())
        self.assertTrue((filter_indexes_by_spacer(d, pad) == np.array([1])).all())
        self.assertTrue((filter_indexes_by_spacer(e, pad) == np.array([1])).all())


if __name__ == "__main__":
    unittest.main()
